
#include <iostream>
#include <vector>
#include <functional>
#include <memory>
#include <exception>
#include <cstring>
#include <uv.h>

class TcpServer
{
public:
    using OnMessageCallback =  std::function<void (const char*, ssize_t)>;

    TcpServer(uv_loop_t *loop): loop(loop)
    {
        memset(&server, 0, sizeof(server));
        client = nullptr;
        port = 0;
        onMessageCallback_ = nullptr;
    }

    ~TcpServer()
    {
    }

    TcpServer(const TcpServer&) = delete;
    TcpServer& operator = (const TcpServer&) = delete;

    int bind_and_listen(uint16_t port)
    {
        this->port = port;
        server.data = this;
        ::uv_tcp_init(loop, &server);
        struct sockaddr_in addr;
        ::uv_ip4_addr("0.0.0.0", port, &addr);
        int r = ::uv_tcp_bind(&server, (const struct sockaddr *) &addr, 0);
        if (r)
        {
            fprintf(stderr, "Bind %d error %s\n", port, uv_strerror(r));
            return -1;
        }
        r = ::uv_listen((uv_stream_t *) &server, 5, on_new_connection);
        if (r) {
            fprintf(stderr, "Listen %d error %s\n", port, uv_strerror(r));
            return -2;
        }
        return 0;
    }

    int close()
    {
        if (client != nullptr)
        {
            ::uv_close((uv_handle_t *)client, client_shutdown_cb);
        }
        return 0;
    }

    void set_recv_callback(OnMessageCallback callback)
    {
        onMessageCallback_ = callback;
    }

    int send(const char *data, ssize_t len)
    {
        if (data && len && client)
        {
            write_req_t *req = (write_req_t*) malloc(sizeof(write_req_t) + len);
            if (!req) return -1;
            req->buf.len = len;
            req->buf.base = (char *)(req + 1);
            memcpy(req->buf.base, data, len);
            return ::uv_write((uv_write_t*) req, (uv_stream_t *)client, &req->buf, 1, write_done);
        }
        return -2;
    }

private:
    typedef struct {
        uv_write_t req;
        uv_buf_t buf;
    } write_req_t;

    static void alloc_buffer(uv_handle_t *handle, size_t suggested_size, uv_buf_t *buf)
    {
        buf->len = suggested_size;
        buf->base = static_cast<char *>(malloc(suggested_size));
    }

    static void client_reject_cb(uv_handle_t *client)
    {
        free(client);
    }

    static void client_close_cb(uv_handle_t *client)
    {
        TcpServer *wrapper_server = static_cast<TcpServer *> (client->data);
        free(client);
        wrapper_server->client = nullptr;
    }

    static void client_shutdown_cb(uv_handle_t *client)
    {
        TcpServer *wrapper_server = static_cast<TcpServer *> (client->data);
        free(client);
        wrapper_server->client = nullptr;
        uv_close((uv_handle_t *)&wrapper_server->server, server_close_cb);
    }

    static void server_close_cb(uv_handle_t *server)
    {
        TcpServer *wrapper_server = static_cast<TcpServer *> (server->data);
        uv_stop(wrapper_server->loop);
    }

    static void write_done(uv_write_t *req, int status)
    {
        if (status < 0)
        {
            fprintf(stderr, "Send error %s\n", uv_strerror(status));
            ::uv_close((uv_handle_t *)req->handle, client_close_cb);
        }
        free(req);
    }

    static void on_new_connection(uv_stream_t *server, int status)
    {
        TcpServer *wrapper_server = static_cast<TcpServer *> (server->data);
        if (status < 0) {
            fprintf(stderr, "New connection error %s\n", uv_strerror(status));
            return;
        }

        auto client = (uv_tcp_t *)malloc(sizeof(uv_tcp_t));
        if (!client)
        {
            return;
        }
    
        ::uv_tcp_init(wrapper_server->loop, client);
        client->data = wrapper_server;
        ::uv_accept(server, (uv_stream_t *)client);
        if (wrapper_server->client != nullptr)
        {
            fprintf(stderr, "already exist a connection!\n");
            ::uv_close((uv_handle_t *)client, client_reject_cb);
        }
        else
        {
            wrapper_server->client = client;
            ::uv_read_start((uv_stream_t *)client, alloc_buffer, read_cb);
            fprintf(stderr, "a new connection!\n");
        }
    }

    static void read_cb(uv_stream_t *client, ssize_t nread, const uv_buf_t *buf)
    {
        TcpServer *wrapper_server = static_cast<TcpServer *> (client->data);
        if (nread > 0)
        {
            wrapper_server->onMessageCallback_(buf->base, nread);
        }
        else if (nread < 0)
        {
            if (nread != UV_EOF) {
                fprintf(stderr, "Read error %s\n", uv_err_name(nread));
            } else {
                fprintf(stderr, "Client disconnect\n");
            }
            ::uv_close((uv_handle_t *) client, client_close_cb);
        }
        if (buf->base != NULL) {
            free(buf->base);
        }
    }

    uv_loop_t *loop;
    uv_tcp_t server;
    uv_tcp_t *client;
    uint16_t port;
    OnMessageCallback onMessageCallback_;
};

int main()
{
    TcpServer s(uv_default_loop());
    s.set_recv_callback([&](const char* data, ssize_t len) {
        if(data[0] == 'a')
        {
            s.close();
        }
        else
        {
            s.send(data, len);
        }
    });
    s.bind_and_listen(7000);
    ::uv_run(uv_default_loop(), UV_RUN_DEFAULT);
    uv_loop_close(uv_default_loop());
    return 0;
}
